# models/time_series_data.py
import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler


class TimeSeriesData:
    """
    多变量时间序列数据处理工具
    支持从长格式数据加载、转换为宽格式、标准化、时间特征提取、数据集构建
    """

    def __init__(self, data_path, target_metrics=None, seq_len=96, pred_len=24):
        """
        初始化数据处理器

        参数:
        - data_path: str, CSV 文件路径
        - target_metrics: list, 要预测的指标列表
        - seq_len: int, 历史序列长度
        - pred_len: int, 预测序列长度
        """
        self.target_metrics = target_metrics or ["usage_active", "mem_used", "disk_io", "network_in", "network_out"]
        self.seq_len = seq_len
        self.pred_len = pred_len
        self.scaler = StandardScaler()
        self.load_data(data_path)

    def load_data(self, data_path):
        """
        加载并预处理原始数据（支持长格式）

        参数:
        - data_path: str, CSV 文件路径
        """
        # 读取原始数据
        df = pd.read_csv(data_path)

        # 检查是否为长格式数据（使用你提供的字段名）
        if 'metric_name' in df.columns and 'value' in df.columns and 'timestamp_value' in df.columns:
            # 转换为宽格式
            df = df.pivot(index='timestamp_value', columns='metric_name', values='value').reset_index()

        # 检查时间戳列并设置索引
        if 'timestamp_value' in df.columns:
            df['timestamp'] = pd.to_datetime(df['timestamp_value'])
            df.set_index('timestamp', inplace=True)
        elif 'timestamp' in df.columns:
            df['timestamp'] = pd.to_datetime(df['timestamp'])
            df.set_index('timestamp', inplace=True)
        else:
            raise KeyError("找不到时间戳列，请确保数据包含 'timestamp' 或 'timestamp_value' 列")

        # 确保目标指标列存在
        missing_cols = [col for col in self.target_metrics if col not in df.columns]
        if missing_cols:
            raise KeyError(f"以下列不存在于 CSV 文件中: {missing_cols}")

        # 重采样为每分钟数据
        df = df[self.target_metrics].reindex(pd.date_range(
            start=df.index.min(), end=df.index.max(), freq='1min'
        )).ffill()

        # 时间戳处理
        df.index = pd.to_datetime(df.index)
        self.ts = df

        # 时间特征
        self.time_features = np.stack([
            self.ts.index.hour.values,
            self.ts.index.dayofweek.values
        ], axis=-1)

        # 数据标准化
        self.ts_values = self.scaler.fit_transform(self.ts.values)

    def create_dataset(self):
        """
        创建训练数据集

        返回:
        - X: 历史序列 [num_samples, seq_len, n_vars]
        - time_X: 时间特征 [num_samples, seq_len, 2]
        - y: 预测目标 [num_samples, pred_len]
        """
        X, time_X, y = [], [], []

        for i in range(len(self.ts_values) - self.seq_len - self.pred_len):
            # 历史窗口
            X.append(self.ts_values[i:i + self.seq_len])
            time_X.append(self.time_features[i:i + self.seq_len])

            # 预测目标（使用第一个指标作为预测目标）
            y.append(self.ts_values[i + self.seq_len:i + self.seq_len + self.pred_len, 0])

        return np.array(X), np.array(time_X), np.array(y).reshape(-1, self.pred_len)
